# coding:utf-8
import pymysql
# import sqlite3
from retrying import retry
from utils.logger import Logger
# import psycopg2
from configs.config import env
import os

os.environ['NLS_LANG'] = 'SIMPLIFIED CHINESE_CHINA.UTF8'


class DBHandle:
    def __init__(self, whichDB, auto_close=True):
        self.DB = env['db'][whichDB]
        self.auto_close = auto_close
        self.connection = None

    @retry(stop_max_attempt_number=10)
    def __connect(self):
        try:
            if self.DB['TYPE'].lower() == 'mysql':
                self.conn = pymysql.connect(**self.DB['DBINFO'])
                self.cur = self.conn.cursor()
            elif self.DB['TYPE'] == 'psycopg2':
                self.conn = psycopg2.connect(**self.DB['DBINFO'])
                self.cur = self.conn.cursor()
            elif self.DB['TYPE'].lower() == 'oracle':
                self.conn = cx_Oracle.connect(self.DB['DBINFO'])
                self.cur = self.conn.cursor()
            # elif self.DB['TYPE'].lower() == 'sqlite':
            #     self.conn = sqlite3.connect(self.DB['DBINFO'])
            #     self.cur = self.conn.cursor()
            else:
                Logger.info('%s is Unsupported Type' % self.DB['TYPE'])
        except Exception as err:
            Logger.error('数据库接连出错, 错误信息:\n%s' % err)
            assert False

    def insert(self, table, data):
        """
        mysql insert() function
        """
        if not self.connection:
            self.__connect()
        with self.cur as cursor:
            params = self.join_field_value(data);
            sql = "INSERT IGNORE INTO {table} SET {params}".format(table=table, params=params)

            result = cursor.execute(sql, tuple(data.values()))
            self.conn.commit()
            return result

    def delete(self, table, condition=None, limit=None):
        """
        mysql delete() function
        sql.PreparedStatement method
        """
        if not self.connection:
            self.__connect()
        with self.cur as cursor:
            prepared = []  # PreparedStatement
            if not condition:
                where = '1';
            elif isinstance(condition, dict):
                where = self.join_field_value(condition, ' AND ')
                prepared.extend(condition.values())
            else:
                where = condition

            limits = "LIMIT {limit}".format(limit=limit) if limit else ""
            sql = "DELETE FROM {table} WHERE {where} {limits}".format(
                table=table, where=where, limits=limits)

            # check PreparedStatement
            if not prepared:
                result = cursor.execute(sql)
            else:
                result = cursor.execute(sql, tuple(prepared))

            self.conn.commit()  # not autocommit

            return result

    def update(self, table, data, condition=None):
        """
        mysql update() function
        Use sql.PreparedStatement method
        """
        if not self.connection:
            self.__connect()
        with self.cur as cursor:
            prepared = [] # PreparedStatement
            params = self.join_field_value(data)
            prepared.extend(data.values())
            if not condition:
                where = '1';
            elif isinstance(condition, dict):
                where = self.join_field_value( condition, ' AND ' )
                prepared.extend(condition.values())
            else:
                where = condition

            sql = "UPDATE {table} SET {params} WHERE {where}".format(
                table=table, params=params, where=where)

            # check PreparedStatement
            if not prepared:
                result = cursor.execute(sql)
            else:
                result = cursor.execute(sql, tuple(prepared))

            self.conn.commit() # not autocommit
            return result

    def count(self, table, condition=None):
        """
        count database record
        """
        if not self.connection:
            self.__connect()
        with self.cur as cursor:
            prepared = []  # PreparedStatement

            # WHERE CONDITION
            if not condition:
                where = '1';
            elif isinstance(condition, dict):
                where = self.join_field_value(condition, ' AND ')
                prepared.extend(condition.values())
            else:
                where = condition

            # SELECT COUNT(*) as cnt
            sql = "SELECT COUNT(*) as cnt FROM {table} WHERE {where}".format(
                table=table, where=where)

            # check PreparedStatement, EXECUTE SELECT COUNT sql
            if not prepared:
                cursor.execute(sql)
            else:
                cursor.execute(sql, tuple(prepared))
            Logger.info(sql)
            # RETURN cnt RESULT
            return cursor.fetchone()[0]

    def fetch_rows(self, table, fields=None, condition=None, order=None, limit=None, fetchone=False):
        """
        mysql select() function
        """
        if not self.connection:
            self.__connect()
        with self.cur as cursor:
            prepared = []  # PreparedStatement

            # SELECT FIELDS
            if not fields:
                fields = '*'
            elif isinstance(fields, tuple) or isinstance(fields, list):
                fields = '`, `'.join(fields)
                fields = '`{fields}`'.format(fields=fields)
            else:
                fields = fields

            # WHERE CONDITION
            if not condition:
                where = '1';
            elif isinstance(condition, dict):
                where = self.join_field_value(condition, ' AND ')
                prepared.extend(condition.values())
            else:
                where = condition

            # ORDER BY OPTIONS
            if not order:
                orderby = ''
            else:
                orderby = 'ORDER BY {order}'.format(order=order)

            # LIMIT NUMS
            limits = "LIMIT {limit}".format(limit=limit) if limit else ""
            sql = "SELECT {fields} FROM {table} WHERE {where} {orderby} {limits}".format(
                fields=fields, table=table, where=where, orderby=orderby, limits=limits)

            # check PreparedStatement
            if not prepared:
                cursor.execute(sql)
            else:
                cursor.execute(sql, tuple(prepared))

            if fetchone:
                return cursor.fetchone()
            else:
                field = [i[0] for i in self.cur.description]
                data = cursor.fetchall()
                n = len(data)
                Logger.info(sql)
                Logger.info('查询出%s条记录' % n)
                result = [{field[f].lower(): data[i][f] for f in range(len(field))} for i in range(n)]
                return result

    def query(self, sql, fetchone=False):
        """
        execute custom sql query
        """
        if not self.connection:
            self.__connect()
        with self.cur as cursor:
            if not sql:
                return
            cursor.execute(sql)
            self.conn.commit() # not auto commit
            if fetchone:
                return cursor.fetchone()
            else:
                return cursor.fetchall()

    def __del__(self):
        """close mysql database connection"""
        self.close()

    def join_field_value(self, data, glue = ', '):
        sql = comma = ''
        for key in data.keys():
            sql +=  "{}`{}` = %s".format(comma, key)
            comma = glue
        return sql

    def execute(self, sql, params=None):
        sql = sql.strip()
        if sql.split()[0].lower() == 'select':
            sqlType = 1
        elif sql.split()[0].lower() in ('update', 'delete', 'insert'):
            sqlType = 2
        else:
            sqlType = 3
        if sqlType == 1:
            if not self.connection:
                self.__connect()
            Logger.info('执行SQL: %s, 参数params为: %s' % (sql,params))
            try:
                if params:
                    self.cur.execute(sql, params)
                else:
                    self.cur.execute(sql)
                field = [i[0] for i in self.cur.description]
                data = self.cur.fetchall()
                n = len(data)
                Logger.info('查询出%s条记录' % n)

                result = [{field[f].lower(): data[i][f] for f in range(len(field))} for i in range(n)]
                if self.auto_close:
                    self.cur.close()
                    self.conn.close()
                return result
            except Exception as err:
                Logger.error('查询SQL出错, 错误信息:\n%s' % err)
                if self.auto_close:
                    self.cur.close()
                    self.conn.close()
                assert False
        if sqlType == 2:
            if not self.connection:
                self.__connect()
            Logger.info('执行SQL: %s' % sql)
            try:
                if params:
                    self.cur.execute(sql, params)
                else:
                    self.cur.execute(sql)
                self.conn.commit()
                if self.auto_close:
                    self.cur.close()
                    self.conn.close()
            except Exception as  err:
                Logger.error('执行SQL出错, 错误信息:\n%s' % err)
                self.conn.rollback()
                if self.auto_close:
                    self.cur.close()
                    self.conn.close()
                assert False
        if sqlType == 3:
            Logger.error('不支持的操作,或sql有误:\n%s' % sql)
            assert False

    def executemany(self, sql, params=None):
        if not self.connection:
            self.__connect()
        try:
            self.cur.executemany(sql, params)
            self.conn.commit()
            if self.auto_close:
                self.cur.close()
                self.conn.close()
        except Exception as err:
            Logger.error('执行SQL出错, 错误信息:\n' % err)
            self.conn.rollback()
            if self.auto_close:
                self.cur.close()
                self.conn.close()
            assert False

    def close(self):
        if self.connection:
            self.cur.close()
            self.conn.close()